import os
import numpy as np
import cv2 as cv
import itertools
from .method_base import CornerDetectMethodBase
from functools import cached_property


class GradientBase(CornerDetectMethodBase):
    def run(
        self,
        output_path: os.PathLike,
        ksize: int,
        sigma: int,
        wsize: int,
        interp: str = "opencv",
        filter: str = "vicinity",
        thres: float = 5e-2,
        *args,
        **kwargs
    ):
        x0s, y0s = self._process(ksize=ksize, wsize=wsize, sigma=sigma, *args, **kwargs)
        img_resized = None
        for x0, y0 in zip(x0s, y0s):
            try:
                x0, y0 = self.filter_method(
                    filter=filter,
                    x0=x0,
                    y0=y0,
                    wsize=wsize,
                    thres=thres,
                    *args,
                    **kwargs
                )
            except:
                continue
            try:
                x0, y0 = self.fit_method(
                    interp=interp,
                    x0=x0,
                    y0=y0,
                    wsize=wsize,
                    sigma=sigma,
                    ksize=ksize,
                    *args,
                    **kwargs
                )
            except Exception as e:
                continue
            if x0 != None and y0 != None:
                img_resized = super().run(
                    output_path=output_path,
                    ksize=ksize,
                    sigma=sigma,
                    x0=x0,
                    y0=y0,
                    wsize=wsize,
                    img_resized=img_resized,
                    *args,
                    **kwargs
                )

    def none_method(self, x0, y0, *args, **kwargs):
        return x0, y0

    def filter_method(self, filter: str, *args, **kwargs):
        filter_methods = {"vicinity": self._filter_vicinity, "none": self.none_method}
        return filter_methods[filter](*args, **kwargs)

    def fit_method(self, interp: str, *args, **kwargs):
        interp_methods = {
            "gauss_fit": self._interp_gauss_fit,
            "cross_line": self._interp_cross_line,
            "opencv": self._interp_opencv,
            "taylor": self._interp_taylor,
            "none": self.none_method,
        }
        return interp_methods[interp](*args, **kwargs)

    def set_gauss(self, ksize=5, sigma=1.5, *args, **kwargs):
        self.ksize = ksize
        self.sigma = sigma

    @cached_property
    def _gx(self):
        return cv.Sobel(
            cv.GaussianBlur(
                self.fig,
                ksize=(self.ksize, self.ksize),
                sigmaX=self.sigma,
                sigmaY=self.sigma,
            ),
            cv.CV_64F,
            1,
            0,
            ksize=self.ksize,
        )

    @cached_property
    def _gy(self):
        return cv.Sobel(
            cv.GaussianBlur(
                self.fig,
                ksize=(self.ksize, self.ksize),
                sigmaX=self.sigma,
                sigmaY=self.sigma,
            ),
            cv.CV_64F,
            0,
            1,
            ksize=self.ksize,
        )

    @cached_property
    def _gxx(self):
        return cv.Sobel(
            cv.GaussianBlur(
                self.fig,
                ksize=(self.ksize, self.ksize),
                sigmaX=self.sigma,
                sigmaY=self.sigma,
            ),
            cv.CV_64F,
            2,
            0,
            ksize=self.ksize,
        )

    @cached_property
    def _gyy(self):
        return cv.Sobel(
            cv.GaussianBlur(
                self.fig,
                ksize=(self.ksize, self.ksize),
                sigmaX=self.sigma,
                sigmaY=self.sigma,
            ),
            cv.CV_64F,
            0,
            2,
            ksize=self.ksize,
        )

    @cached_property
    def _gxy(self):
        return cv.Sobel(
            cv.GaussianBlur(
                self.fig,
                ksize=(self.ksize, self.ksize),
                sigmaX=self.sigma,
                sigmaY=self.sigma,
            ),
            cv.CV_64F,
            1,
            1,
            ksize=self.ksize,
        )

    def _filter_vicinity(self, x0, y0, thres, wsize, *args, **kwargs):
        for dx in range(-wsize // 2, wsize // 2 + 1):
            for dy in range(-wsize // 2, 1):
                if dx == 0 and dy == 0:
                    continue
                if (
                    np.abs(
                        (self._gy[x0 + dx, y0 + dy] - self._gy[x0 - dx, y0 - dy])
                        / (1 + self._gx[x0 + dx, y0 + dy] * self._gx[x0 - dx, y0 - dy])
                    )
                    > thres
                ):
                    return x0, y0
        return None, None

    def _interp_gauss_fit(self, x0, y0, wsize, *args, **kwargs):
        """
        高斯曲面拟合法，没有搞明白
        """
        # 生成横坐标序列
        row_axis = (
            np.ones([wsize, 1], dtype=np.uint32)
            @ np.arange(x0, wsize + x0, dtype=np.uint32).reshape([1, wsize])
            + x0
        )
        # 生成纵坐标序列
        col_axis = (
            np.arange(y0, wsize + y0, dtype=np.uint32).reshape([wsize, 1])
            @ np.ones([1, wsize], dtype=np.uint32)
            + y0
        )
        X = np.concatenate(
            (
                (row_axis**2).reshape([1, wsize * wsize]),
                (col_axis**2).reshape([1, wsize * wsize]),
                row_axis.reshape([1, wsize * wsize]),
                col_axis.reshape([1, wsize * wsize]),
                np.ones_like(col_axis.reshape([1, wsize * wsize])),
            )
        )
        Y = self.fig[
            x0 - wsize // 2 : x0 + wsize // 2 + 1, y0 - wsize // 2 : y0 + wsize // 2 + 1
        ].reshape([1, wsize * wsize])
        ind = Y > 0
        Y = Y[ind]
        X = X[:, ind[0]]
        A = (np.log(Y) @ X.T @ np.matrix(X @ X.T).I).tolist()
        A = A[0]
        x_0 = -A[2] / 2 / A[0]
        y_0 = -A[3] / 2 / A[1]
        return x_0, y_0

    def _interp_taylor(self, x0, y0, wsize, *args, **kwargs):
        """
        多元函数泰勒展式拟合
        """
        try:
            denominator = self._gxx[x0, y0] * self._gyy[x0, y0] - self._gxy[x0, y0] ** 2
        except:
            raise ValueError
        s = (
            self._gy[x0, y0] * self._gxy[x0, y0] - self._gx[x0, y0] * self._gyy[x0, y0]
        ) / (denominator + 1e-9)
        t = (
            self._gx[x0, y0] * self._gxy[x0, y0] - self._gy[x0, y0] * self._gxx[x0, y0]
        ) / (denominator + 1e-9)
        if s**2 + t**2 < 2:
            return x0 + s, y0 + t
        x0 = int(x0 + s)
        y0 = int(x0 + t)
        return None, None

    def _interp_cross_line(self, x0, y0, wsize, ksize, *args, **kwargs):
        """
        边缘相交法拟合角点，首先在横纵方向上检测一波线条中心
        """
        wsize = 18

        def __detect_line(x0, y0, d_fig, ksize, *args, **kwargs):
            temp_fig = d_fig**2
            temp_fig /= np.max(temp_fig)
            fig = np.argwhere(temp_fig > 0.9)
            x, y = fig[:, 0] + x0 - wsize // 2, fig[:, 1] + y0 - wsize // 2
            for x0, y0 in zip(x, y):
                cv.drawMarker(
                    temp_fig, (int(y0), int(x0)), (0, 0, 255), cv.MARKER_CROSS, 10, 1
                )
            cv.imwrite(
                "test.png",
                (temp_fig - np.min(temp_fig))
                / (np.max(temp_fig) - np.min(temp_fig))
                * 255,
            )
            return fig[:, 0] + x0 - wsize // 2, fig[:, 1] + y0 - wsize // 2

        temp_fig = self.fig[
            x0 - wsize // 2 : x0 + wsize // 2 + 1, y0 - wsize // 2 : y0 + wsize // 2 + 1
        ]
        d_fig = cv.Sobel(temp_fig, cv.CV_64F, 1, 0, ksize=3)
        # 生成横坐标序列
        row_axis, col_axis = __detect_line(x0, y0, d_fig, ksize, *args, **kwargs)
        len_row_axis = len(row_axis)
        X = np.concatenate(
            [
                row_axis.reshape([1, len_row_axis]),
                col_axis.reshape([1, len_row_axis]),
                np.ones_like(col_axis).reshape([1, len_row_axis]),
            ],
        )
        eigenvalues, eigenvectors = np.linalg.eig(np.matrix(X @ X.T))
        min_eigenvalue_index = np.argmin(np.abs(eigenvalues))
        A1 = eigenvectors[:, min_eigenvalue_index]

        # temp_fig = self._gy[
        #     x0 - wsize // 2 : x0 + wsize // 2 + 1, y0 - wsize // 2 : y0 + wsize // 2 + 1
        # ]
        d_fig = cv.Sobel(temp_fig, cv.CV_64F, 0, 1, ksize=3)
        row_axis, col_axis = __detect_line(x0, y0, d_fig, ksize, *args, **kwargs)
        len_row_axis = len(row_axis)
        X_ = np.concatenate(
            (
                row_axis.reshape([1, len_row_axis]),
                col_axis.reshape([1, len_row_axis]),
                np.ones_like(col_axis).reshape([1, len_row_axis]),
            )
        )
        eigenvalues, eigenvectors = np.linalg.eig(np.matrix(X_ @ X_.T))
        min_eigenvalue_index = np.argmin(np.abs(eigenvalues))
        A2 = eigenvectors[:, min_eigenvalue_index]
        Y = np.cross(A1.T, A2.T)
        return Y[0, 0] / Y[0, 2], Y[0, 1] / Y[0, 2]

    def _interp_opencv(self, x0, y0, wsize, *args, **kwargs):
        """
        opencv 自带工具箱拟合
        """
        winSize = (wsize, wsize)
        zeroZone = (-1, -1)
        criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 100, 0.001)
        x0_y0 = cv.cornerSubPix(
            image=self.fig,
            corners=np.array([[y0, x0]], dtype=np.float32),
            winSize=winSize,
            zeroZone=zeroZone,
            criteria=criteria,
        )
        if (x0_y0[0, 0] - y0) ** 2 + (x0_y0[0, 1] - x0) ** 2 > 4:
            return None, None
        return x0_y0[0, 1], x0_y0[0, 0]


class KitchenRosenfeldMethod(GradientBase):
    @property
    def _name(self):
        return "KitchenRosenfeldMethod"

    def _process(self, *args, **kwargs):
        self.set_gauss(*args, **kwargs)
        denominator = self._gx**2 + self._gy**2
        nominator = (
            self._gxx * self._gy**2
            + self._gyy * self._gx**2
            - 2 * self._gxy * self._gx * self._gy
        )
        k = np.abs(nominator / (denominator + 1e-9))
        k /= np.max(k)
        x0_y0 = np.argwhere(k > 0.8)
        try:
            return x0_y0[:, 0], x0_y0[:, 1]
        except:
            return x0_y0[0, 0], x0_y0[0, 1]


class HarrisMethod(GradientBase):
    @property
    def _name(self):
        return "HarrisMethod"

    def _process(self, ksize, sigma, *args, **kwargs):
        self.set_gauss(ksize, sigma, *args, **kwargs)
        A = cv.GaussianBlur(self._gxx, (ksize, ksize), sigma)
        B = cv.GaussianBlur(self._gyy, (ksize, ksize), sigma)
        C = cv.GaussianBlur(self._gxy, (ksize, ksize), sigma)
        R = np.abs(A * B - C**2 - 0.04 * (A + B) ** 2)
        R /= np.max(R) + 1e-9
        x0_y0 = np.argwhere(R > 0.8)
        return x0_y0[:, 0], x0_y0[:, 1]


class OpencvMethod(GradientBase):
    @property
    def _name(self):
        return "Opencv"

    def _process(self, ksize, wsize, sigma, *args, **kwargs):
        self.set_gauss(*args, **kwargs)
        R = np.abs(cv.cornerHarris(self.fig, blockSize=wsize, ksize=ksize, k=0.04))
        R /= np.max(R) + 1e-9
        x0_y0 = np.argwhere(R > 0.8)
        return x0_y0[:, 0], x0_y0[:, 1]


class HeissanMethod(GradientBase):
    @property
    def _name(self):
        return "HeissanMethod"

    def _process(self, ksize: int, sigma: int, *args, **kwargs):
        self.set_gauss(*args, **kwargs)
        C = np.abs(self._gxx * self._gyy + self._gxy**2)
        if (C - np.abs(self._gxx - self._gyy)).all() == 0:
            C /= np.max(C) + 1e-9
            x0_y0 = np.argwhere(C > 0.9)
            return x0_y0[:, 0], x0_y0[:, 1]
        else:
            raise Exception("错误！两个C值阵不相等！")
